import numpy as np
import torch
import torch.utils.data as Data
from nndl.dataset import load_data
import torch.nn as nn


class IrisDataset(Data.Dataset):
    def __init__(self, mode='train', num_train=120, num_dev=15):
        super(IrisDataset, self).__init__()
        # 调用第三章中的数据读取函数，其中不需要将标签转成one-hot类型
        X, y = load_data(shuffle=True)
        if mode == 'train':
            self.X, self.y = X[:num_train], y[:num_train]
        elif mode == 'dev':
            self.X, self.y = X[num_train:num_train + num_dev], y[num_train:num_train + num_dev]
        else:
            self.X, self.y = X[num_train + num_dev:], y[num_train + num_dev:]

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

    def __len__(self):
        return len(self.y)


train_dataset = IrisDataset(mode='train')
dev_dataset = IrisDataset(mode='dev')
test_dataset = IrisDataset(mode='test')

print("length of train set: ", len(train_dataset))

batch_size = 16

# 加载数据
train_loader = Data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
dev_loader = Data.DataLoader(dev_dataset, batch_size=batch_size)
test_loader = Data.DataLoader(test_dataset, batch_size=batch_size)

from sklearn.svm import SVC
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import accuracy_score
def SVC_split(x_train, y_train, x_test, y_test):
    # 定义SVM分类器
    svm = SVC()

    # 定义参数空间
    param_grid = {
        'C': [0.01, 0.1, 1, 10],
        'kernel': ['linear', 'rbf', 'poly'],
        'gamma': ['scale', 'auto']
    }
    # 使用GridSearchCV进行交叉验证和参数选择
    grid_search = GridSearchCV(svm, param_grid=param_grid, cv=5)
    grid_search.fit(x_train, y_train.ravel())

    # 在测试集上评估模型性能
    y_pred = grid_search.predict(x_test)
    accuracy = accuracy_score(y_test, y_pred)
    print("Accuracy on test set: {:.2f}%".format(accuracy * 100))

# 定义前馈神经网络
class Model_MLP_L2_V3(nn.Module):
    def __init__(self, input_size, output_size, hidden_size):
        super(Model_MLP_L2_V3, self).__init__()
        # 构建第一个全连接层
        self.fc1 = nn.Linear(input_size, hidden_size)
        # 构建第二全连接层
        self.fc2 = nn.Linear(hidden_size, output_size)
        # 定义网络使用的激活函数
        self.act = nn.Sigmoid()
        # nn.init.normal_(self.fc1.weight, mean=0., std=0.01)
        w_ = torch.normal(0, 0.01, size=(hidden_size, input_size), requires_grad=True)
        self.fc1.weight = torch.nn.Parameter(w_)
        nn.init.constant_(self.fc1.bias, 1.0)
        # nn.init.normal_(self.fc2.weight, mean=0., std=0.01)
        w2 = torch.normal(0, 0.01, size=(output_size, hidden_size), requires_grad=True)
        self.fc2.weight = nn.Parameter(w2)
        nn.init.constant_(self.fc2.bias, 1.0)

    def forward(self, inputs):
        outputs = self.fc1(inputs)
        outputs = self.act(outputs)
        outputs = self.fc2(outputs)
        return outputs


fnn_model = Model_MLP_L2_V3(input_size=4, output_size=3, hidden_size=6)


class Accuracy(object):
    def __init__(self, is_logist=True):
        """
        输入：
           - is_logist: outputs是logist还是激活后的值
        """

        # 用于统计正确的样本个数
        self.num_correct = 0
        # 用于统计样本的总数
        self.num_count = 0

        self.is_logist = is_logist

    def update(self, outputs, labels):
        """
        输入：
           - outputs: 预测值, shape=[N,class_num]
           - labels: 标签值, shape=[N,1]
        """

        # 判断是二分类任务还是多分类任务，shape[1]=1时为二分类任务，shape[1]>1时为多分类任务
        if outputs.shape[1] == 1:  # 二分类
            outputs = torch.squeeze(outputs, axis=-1)
            if self.is_logist:
                # logist判断是否大于0
                preds = (outputs >= 0).to(torch.float32)
            else:
                # 如果不是logist，判断每个概率值是否大于0.5，当大于0.5时，类别为1，否则类别为0
                preds = (outputs >= 0.5).to(torch.float32)
        else:
            # 多分类时，使用'paddle.argmax'计算最大元素索引作为类别
            preds = torch.argmax(outputs, dim=1).int()

        # 获取本批数据中预测正确的样本个数
        labels = torch.squeeze(labels, dim=-1)
        # batch_correct = torch.sum(torch.tensor(preds == labels, dtype=torch.float32)).numpy()
        # batch_correct = (preds == labels).sum().item()
        batch_correct = float((preds == labels).sum())
        batch_count = len(labels)

        # 更新num_correct 和 num_count
        self.num_correct += batch_correct
        self.num_count += batch_count

    def accumulate(self):
        # 使用累计的数据，计算总的指标
        if self.num_count == 0:
            return 0
        return self.num_correct / self.num_count

    def reset(self):
        # 重置正确的数目和总数
        self.num_correct = 0
        self.num_count = 0

    def name(self):
        return "Accuracy"


import torch.optim as opt
from nndl.runner import RunnerV3
import torch.nn.functional as F

lr = 0.2
# 定义网络
model = fnn_model
# 定义优化器
optimizer = opt.SGD(model.parameters(), lr=lr)
# 定义损失函数。softmax+交叉熵
loss_fn = F.cross_entropy

# 定义评价指标
metric = Accuracy(is_logist=True)

runner = RunnerV3(model, optimizer, loss_fn, metric)

# 启动训练
log_steps = 100
eval_steps = 50
runner.train(train_loader, dev_loader,
             num_epochs=150, log_steps=log_steps, eval_steps=eval_steps,
             save_path="best_model.pdparams")


import matplotlib.pyplot as plt


# 绘制训练集和验证集的损失变化以及验证集上的准确率变化曲线
def plot_training_loss_acc(runner, fig_name,
                           fig_size=(16, 6),
                           sample_step=20,
                           loss_legend_loc="upper right",
                           acc_legend_loc="lower right",
                           train_color="#e4007f",
                           dev_color='#f19ec2',
                           fontsize='large',
                           train_linestyle="-",
                           dev_linestyle='--'):
    global dev_steps
    plt.figure(figsize=fig_size)

    plt.subplot(1, 2, 1)
    train_items = runner.train_step_losses[::sample_step]
    train_steps = [x[0] for x in train_items]
    train_losses = [x[1] for x in train_items]

    plt.plot(train_steps, train_losses, color=train_color, linestyle=train_linestyle, label="Train loss")
    if len(runner.dev_losses) > 0:
        dev_steps = [x[0] for x in runner.dev_losses]
        dev_losses = [x[1] for x in runner.dev_losses]
        plt.plot(dev_steps, dev_losses, color=dev_color, linestyle=dev_linestyle, label="Dev loss")
    # 绘制坐标轴和图例
    plt.ylabel("loss", fontsize=fontsize)
    plt.xlabel("step", fontsize=fontsize)
    plt.legend(loc=loss_legend_loc, fontsize='x-large')

    # 绘制评价准确率变化曲线
    if len(runner.dev_scores) > 0:
        plt.subplot(1, 2, 2)
        plt.plot(dev_steps, runner.dev_scores,
                 color=dev_color, linestyle=dev_linestyle, label="Dev accuracy")

        # 绘制坐标轴和图例
        plt.ylabel("score", fontsize=fontsize)
        plt.xlabel("step", fontsize=fontsize)
        plt.legend(loc=acc_legend_loc, fontsize='x-large')

    plt.savefig(fig_name)
    plt.show()


plot_training_loss_acc(runner, 'fw-loss.pdf')

# 加载最优模型
runner.load_model('best_model.pdparams')
# 模型评价
score, loss = runner.evaluate(test_loader)
print("[Test] accuracy/loss: {:.4f}/{:.4f}".format(score, loss))

# 获取测试集中第一条数据
X, label = train_dataset[0]
logits = runner.predict(X)


pred_class = torch.argmax(logits[0]).numpy()
label = label.numpy()

# 输出真实类别与预测类别
print("The true category is {} and the predicted category is {}".format(label, pred_class))

SVC_split(train_dataset.X,train_dataset.y ,test_dataset.X, test_dataset.y)